#include "consistent_hashing.hh"
#include "../logger/writer.hh"
#include "../thirdparty/md5/md5.h"
#include <algorithm>

namespace kratos { namespace corelib {

ConsistentHash::ConsistentHash()
: start_(false) {
}

ConsistentHash::~ConsistentHash()  {
}

bool ConsistentHash::start(RehashCallback rehashCallback, StringHashMethod hashMethod) {
    if (start_) {
        return true;
    }
    if (!hashMethod) {
        hashMethod_ = [](const std::string& s) {
            MD5 md5(s);
            std::uint64_t seed = 131;
            std::uint64_t hash = 0;
            for (auto c : md5.md5()) {
                hash = hash * seed + c;
            }
            return (std::uint32_t)(hash & 0x7FFFFFFFFFFFFFFF);
        };
    } else {
        hashMethod_ = hashMethod;
    }
    rehashCallback_ = rehashCallback;
    start_ = true;
    return true;
}

bool ConsistentHash::addNode(const std::string & name, std::uint64_t value, std::size_t count) {
    if (name.empty() || (value == 0) || (count == 0) || (!hashMethod_)) {
        return false;
    }
    for (std::size_t i = 0; i < count; i++) {
        auto key = hashMethod_(name + "#" + std::to_string(i));
        if (!circle_.addNode(key, value, *this)) {
            return false;
        }
    }
    nodeMap_[name] = count;
    return true;
}

bool ConsistentHash::addNode(std::uint64_t key, std::uint64_t value, std::size_t count) {
    return addNode(std::to_string(key), value, count);
}

bool ConsistentHash::removeNode(const std::string & name) {
    if (name.empty() || (!hashMethod_)) {
        return false;
    }
    auto it = nodeMap_.find(name);
    if (it == nodeMap_.end()) {
        return false;
    }
    for (std::size_t i = 0; i < it->second; i++) {
        auto key = hashMethod_(name + "#" + std::to_string(i));
        if (!circle_.removeNode(key, *this)) {
            return false;
        }
    }
    nodeMap_.erase(it);
    return true;
}

bool ConsistentHash::removeNode(std::uint64_t key) {
    return removeNode(std::to_string(key));
}

std::uint64_t ConsistentHash::getNode(std::uint64_t key, HashType type) {
    if (!key || (!hashMethod_)) {
        return 0;
    }
    std::uint64_t value = 0;
    circle_.findNode(key,
        [&](std::uint64_t v) {
            value = v;
        },
        type
    );
    return value;
}

std::size_t ConsistentHash::getVNodeCount(const std::string & name) {
    if (!hashMethod_ || name.empty()) {
        return 0;
    }
    auto it = nodeMap_.find(name);
    if (it == nodeMap_.end()) {
        return 0;
    }
    return it->second;
}

void ConsistentHash::rehash(std::uint64_t oldNode, std::uint64_t newNode) {
    if (rehashCallback_) {
        rehashCallback_(oldNode, newNode);
    }
}

bool HashCircle::addNode(std::uint32_t key, std::uint64_t value, ConsistentHash& hash) {
    for (auto it = circle_.begin(); it != circle_.end(); it++) {
        if (key <= it->key) {
            auto old = it->value;
            circle_.insert(it, HashNode{ key, value });
            hash.rehash(old, value);
            return true;
        }
    }
    circle_.push_back(HashNode{ (std::uint32_t)key, value });
    hash.rehash(0, value);
    return true;
}

bool HashCircle::removeNode(std::uint32_t key, ConsistentHash& hash) {
    for (auto it = circle_.begin(); it != circle_.end(); it++) {
        if (it->key == key) {
            auto old = it->value;
            auto inserter = circle_.erase(it);
            if (inserter == circle_.end()) {
                if (circle_.empty()) {
                    hash.rehash(old, 0);
                } else {
                    hash.rehash(old, circle_.front().value);
                }
            } else {
                hash.rehash(old, inserter->value);
            }
            return true;
        }
    }
    return false;
}

bool HashCircle::findNode(std::uint64_t key, ValueMethod method, HashType type) {
    if (circle_.empty()) {
        return false;
    }
    if (type == HashType::HASH_TYPE_CIRCLE) {
        if (key > circle_.back().key || key < circle_.front().key) {
            method(circle_.front().value);
            return true;
        }
        auto it = std::lower_bound(circle_.begin(), circle_.end(), key, &HashCircle::compare);
        if (it == circle_.end()) {
            return false;
        }
        method(it->value);
    } else if (type == HashType::HASH_TYPE_MOD) {
        auto i = key % (std::uint64_t)circle_.size();
        method(circle_[static_cast<std::size_t>(i)].value);
    }
    return true;
}

}}
